package cn.heycm.blog.common.mybatis;

import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.Collection;
import java.util.function.BiConsumer;

/**
 * 批处理
 *
 * @author hey
 * @version 1.0
 * @since 2024/9/17 22:00
 */
@Component
public class BatchHandler {

    private static final int BATCH_SIZE = 1000;

    @Autowired
    private SqlSessionFactory sqlSessionFactory;

    public <T, M> int execute(Collection<T> list, Class<M> mapperClass, BiConsumer<T, M> consumer) {
        int i = 1;
        SqlSession batchSession = null;
        try {
            batchSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
            M mapper = batchSession.getMapper(mapperClass);
            int size = list.size();
            for (T element : list) {
                consumer.accept(element, mapper);
                if (i % BATCH_SIZE == 0 || i == size) {
                    batchSession.flushStatements();
                }
                i++;
            }
            // 非事务环境下强制commit，事务情况下该commit相当于无效
            batchSession.commit(!TransactionSynchronizationManager.isSynchronizationActive());
        } catch (Exception e) {
            if (batchSession != null) {
                batchSession.rollback();
            }
            throw e;
        } finally {
            if (batchSession != null) {
                batchSession.close();
            }
        }
        return i - 1;
    }
}
